use itertools::{Either, Itertools};
use query_builder::{CreateRecord, CreateRecordDefaultsQuery, QueryBuilder};
use query_core::{
    ConnectRecords, DeleteManyRecords, DeleteRecord, DisconnectRecords, RawQuery, UpdateManyRecords, UpdateRecord,
    UpdateRecordWithSelection, UpdateRecordWithoutSelection, WriteQuery,
};
use query_structure::{
    FieldSelection, Filter, IntoFilter, Model, PrismaValue, PrismaValueType, QueryArguments, RecordFilter,
    RelationLoadStrategy, Take,
};
use sql_query_builder::write::split_write_args_by_shape;
use std::{collections::BTreeMap, iter, mem};

use crate::{
    TranslateError, binding,
    expression::{Binding, Expression, FieldInitializer, FieldOperation, InMemoryOps, Pagination},
    translate::TranslateResult,
};

use super::read::add_inmemory_join;

pub(crate) fn translate_write_query(query: WriteQuery, builder: &dyn QueryBuilder) -> TranslateResult<Expression> {
    Ok(match query {
        WriteQuery::CreateRecord(cr) => {
            let CreateRecord {
                select_defaults,
                insert_query,
                last_insert_id_field,
                merge_values,
            } = builder
                .build_create_record(&cr.model, cr.args, &cr.selected_fields)
                .map_err(TranslateError::QueryBuildFailure)?;

            let mut initializers = merge_values
                .into_iter()
                .map(|(field, value)| (field.db_name().into(), FieldInitializer::Value(value)))
                .collect::<BTreeMap<_, _>>();

            if let Some(last_insert_id_field) = last_insert_id_field {
                initializers.insert(last_insert_id_field.db_name().into(), FieldInitializer::LastInsertId);
            }

            if let Some(defaults_query) = &select_defaults {
                // Extend the result record with any defaults that are generated by the extra
                // select query, since they can a part of the primary identifier.
                for (field, placeholder) in &defaults_query.field_placeholders {
                    let value = FieldInitializer::Value(PrismaValue::Placeholder(placeholder.clone()));
                    initializers.insert(field.db_name().into(), value);
                }
            }

            let mut expr = Expression::Unique(Expression::Query(insert_query).into());

            if !initializers.is_empty() {
                expr = Expression::InitializeRecord {
                    expr: expr.into(),
                    fields: initializers,
                };
            }

            if let Some(CreateRecordDefaultsQuery {
                query,
                field_placeholders,
            }) = select_defaults
            {
                let select_binding =
                    Binding::new(binding::defaults(), Expression::Unique(Expression::Query(query).into()));

                let field_bindings = field_placeholders.into_iter().map(|(field, ph)| {
                    let get_defaults = Expression::Get {
                        name: binding::defaults(),
                    };
                    Binding::new(
                        ph.name,
                        Expression::MapField {
                            field: field.db_name().into(),
                            records: get_defaults.into(),
                        },
                    )
                });

                expr = Expression::Let {
                    bindings: iter::once(select_binding).chain(field_bindings).collect(),
                    expr: expr.into(),
                };
            }

            expr
        }

        WriteQuery::CreateManyRecords(cmr) => {
            let split_args = if cmr.split_by_shape && !cmr.args.is_empty() {
                Either::Left(split_write_args_by_shape(&cmr.model, cmr.args))
            } else {
                Either::Right([cmr.args])
            };
            let (projection, nested) = cmr.selected_fields.map(|sf| (sf.fields, sf.nested)).unzip();

            let inserts = split_args
                .into_iter()
                .map(|args| {
                    builder
                        .build_inserts(&cmr.model, args, cmr.skip_duplicates, projection.as_ref())
                        .map_err(TranslateError::QueryBuildFailure)
                })
                .flatten_ok();

            if projection.is_some() {
                let mut expr = Expression::Concat(inserts.map_ok(Expression::Query).try_collect()?);
                let nested = nested.unwrap_or_default();
                if !nested.is_empty() {
                    expr = add_inmemory_join(expr, nested, builder)?;
                }
                expr
            } else {
                Expression::Sum(inserts.map_ok(Expression::Execute).try_collect()?)
            }
        }

        WriteQuery::UpdateManyRecords(UpdateManyRecords {
            name: _,
            model,
            mut record_filter,
            args,
            selected_fields,
            limit,
        }) => {
            let projection = selected_fields.as_ref().map(|f| &f.fields);

            if args.is_empty() {
                let projection = projection.cloned().unwrap_or_else(|| model.primary_identifier());
                let query = build_query_from_record_filter(builder, &model, record_filter, &projection, None)?;
                return if selected_fields.is_some() {
                    Ok(Expression::Query(query))
                } else {
                    Ok(Expression::Execute(query))
                };
            }

            let selector_bindings = limit
                .map(|limit| extract_selectors_that_require_limit(&mut record_filter, limit))
                .unwrap_or_default();

            let updates = builder
                .build_updates(&model, record_filter, args, projection, limit)
                .map_err(TranslateError::QueryBuildFailure)?
                .into_iter()
                .map(if projection.is_some() {
                    Expression::Query
                } else {
                    Expression::Execute
                })
                .collect::<Vec<_>>();

            let mut expr = if let Some(selected_fields) = selected_fields {
                let mut expr = Expression::Concat(updates);
                if !selected_fields.nested.is_empty() {
                    expr = add_inmemory_join(expr, selected_fields.nested, builder)?;
                }
                expr
            } else {
                Expression::Sum(updates)
            };

            if !selector_bindings.is_empty() {
                expr = Expression::Let {
                    bindings: selector_bindings,
                    expr: expr.into(),
                };
            }

            expr
        }

        WriteQuery::UpdateRecord(UpdateRecord::WithSelection(UpdateRecordWithSelection {
            name: _,
            model,
            record_filter,
            args,
            selected_fields,
            selection_order: _,
        })) => {
            let query = if args.is_empty() {
                // We can just issue a read query if there are no write arguments.
                build_query_from_record_filter(builder, &model, record_filter, &selected_fields, Some(Take::Some(1)))?
            } else {
                builder
                    .build_update(&model, record_filter, args, Some(&selected_fields))
                    .map_err(TranslateError::QueryBuildFailure)?
            };
            Expression::Unique(Box::new(Expression::Query(query)))
        }

        // This case is used for databases with no support for the `RETURNING` clause.
        WriteQuery::UpdateRecord(UpdateRecord::WithoutSelection(UpdateRecordWithoutSelection {
            model,
            record_filter,
            args,
        })) => {
            let id_fields = model.shard_aware_primary_identifier();

            // Initialize the fields of a record that represents the identifier of the row that
            // is being updated. We use the record selectors to populate their values.
            let initializers = record_filter
                .selectors
                .as_ref()
                .expect("should have selectors for update")
                .iter()
                .exactly_one()
                .expect("should have exactly one selector")
                .pairs
                .iter()
                .map(|(field, val)| (field.db_name().into(), FieldInitializer::Value(val.clone())))
                .collect();

            // Keep track of the operations that are applied to the primary identifier or shard key fields.
            // They need to be applied in-memory to the record we intend to return.
            let operations = id_fields
                .selections()
                .filter_map(|field| {
                    Some((
                        field.db_name().into(),
                        FieldOperation::try_from(args.get_field_value(&field.db_name())?.as_scalar()?.clone())
                            .expect("arg value should be convertible to FieldOperation"),
                    ))
                })
                .collect::<BTreeMap<_, _>>();

            let expr = if args.is_empty() {
                // We can just skip the update if there are no write arguments, but the query
                // graph still relies on this being wrapped with an `InitializeRecord`
                // containing the primary identifier, even if no update is performed.
                Expression::Unit
            } else {
                Expression::Execute(
                    builder
                        .build_update(&model, record_filter, args, None)
                        .map_err(TranslateError::QueryBuildFailure)?,
                )
            };

            let mut expr = Expression::InitializeRecord {
                expr: expr.into(),
                fields: initializers,
            };

            if !operations.is_empty() {
                expr = Expression::MapRecord {
                    expr: expr.into(),
                    fields: operations,
                };
            }

            expr
        }

        WriteQuery::Upsert(upsert) => {
            let query = builder
                .build_upsert(
                    upsert.model(),
                    upsert.filter().clone(),
                    upsert.create().clone(),
                    upsert.update().clone(),
                    upsert.selected_fields(),
                    &upsert.unique_constraints(),
                )
                .map_err(TranslateError::QueryBuildFailure)?;
            Expression::Unique(Box::new(Expression::Query(query)))
        }

        WriteQuery::QueryRaw(RawQuery {
            model,
            inputs,
            query_type,
        }) => Expression::Query(
            builder
                .build_raw(model.as_ref(), inputs, query_type)
                .map_err(TranslateError::QueryBuildFailure)?,
        ),

        WriteQuery::ExecuteRaw(RawQuery {
            model,
            inputs,
            query_type,
        }) => Expression::Execute(
            builder
                .build_raw(model.as_ref(), inputs, query_type)
                .map_err(TranslateError::QueryBuildFailure)?,
        ),

        WriteQuery::DeleteRecord(DeleteRecord {
            name: _,
            model,
            record_filter,
            selected_fields,
        }) => {
            let selected_fields = selected_fields.as_ref().map(|sf| &sf.fields);
            let query = builder
                .build_delete(&model, record_filter, selected_fields)
                .map_err(TranslateError::QueryBuildFailure)?;
            if selected_fields.is_some() {
                Expression::Unique(Box::new(Expression::Query(query)))
            } else {
                Expression::Execute(query)
            }
        }

        WriteQuery::DeleteManyRecords(DeleteManyRecords {
            model,
            mut record_filter,
            limit,
        }) => {
            let selector_bindings = limit
                .map(|limit| extract_selectors_that_require_limit(&mut record_filter, limit))
                .unwrap_or_default();

            let mut expr = Expression::Sum(
                builder
                    .build_deletes(&model, record_filter, limit)
                    .map_err(TranslateError::QueryBuildFailure)?
                    .into_iter()
                    .map(Expression::Execute)
                    .collect::<Vec<_>>(),
            );

            if !selector_bindings.is_empty() {
                expr = Expression::Let {
                    bindings: selector_bindings,
                    expr: expr.into(),
                };
            }

            expr
        }

        WriteQuery::ConnectRecords(ConnectRecords {
            parent_id,
            child_ids,
            relation_field,
        }) => {
            let (_, parent) = parent_id
                .into_iter()
                .flat_map(IntoIterator::into_iter)
                .exactly_one()
                .expect("query compiler connects should never have more than one parent expression");
            let (_, child) = child_ids
                .into_iter()
                .flat_map(IntoIterator::into_iter)
                .exactly_one()
                .expect("query compiler connects should never have more than one child expression");
            let query = builder
                .build_m2m_connect(relation_field, parent, child)
                .map_err(TranslateError::QueryBuildFailure)?;
            Expression::Execute(query)
        }

        WriteQuery::DisconnectRecords(DisconnectRecords {
            parent_id,
            child_ids,
            relation_field,
        }) => {
            let parent_id = parent_id.as_ref().expect("should have parent ID for disconnect");
            let query = builder
                .build_m2m_disconnect(relation_field, parent_id, &child_ids)
                .map_err(TranslateError::QueryBuildFailure)?;
            Expression::Execute(query)
        }
    })
}

fn build_query_from_record_filter(
    builder: &dyn QueryBuilder,
    model: &Model,
    record_filter: RecordFilter,
    projection: &FieldSelection,
    take: Option<Take>,
) -> Result<query_builder::DbQuery, TranslateError> {
    // It's possible for us to receive selectors here and they cannot be passed to
    // a read query directly, so we need to build a filter from them.
    let filter = record_filter
        .selectors
        .map(IntoFilter::filter)
        .into_iter()
        .fold(record_filter.filter, |acc, f| Filter::and(vec![acc, f]));
    let mut query_args = QueryArguments::from((model.clone(), filter));
    if let Some(take) = take {
        query_args.take = take;
    }

    let query = builder
        .build_get_records(model, query_args, projection, RelationLoadStrategy::Query)
        .map_err(TranslateError::QueryBuildFailure)?
        .into_iter()
        .exactly_one()
        .expect("should have exactly one query for update with selection");
    Ok(query)
}

/// Extracts selectors from the filter that require an in-memory limit operation as bindings.
/// Selectors in the [`RecordFilter`] are replaced with placeholders that refer to the
/// returned bindings.
fn extract_selectors_that_require_limit(record_filter: &mut RecordFilter, limit: usize) -> Vec<Binding> {
    record_filter
        .selectors
        .iter_mut()
        .flatten()
        .flat_map(|result| result.pairs.iter_mut())
        .filter_map(|(field, value)| {
            let typ = value.r#type();
            if !matches!(typ, PrismaValueType::List(_)) {
                return None;
            }

            let name = binding::selector(field);
            let value = mem::replace(value, PrismaValue::placeholder(name.clone(), typ));
            let pagination = Pagination::builder().take(limit as i64).build();
            let expr = Expression::Value(value);
            Some(Binding::new(name, InMemoryOps::from(pagination).into_expression(expr)))
        })
        .collect_vec()
}
